import numpy as np
import matplotlib.pyplot as plt
from SteerFunction import clothoid_path
import math
import time
import os
def perpendicular_line_segment_params(A, B, C, length):
    """
    Calculate the parameters of a line segment perpendicular to line AB, centered at point C, 
    with a given length.

    Parameters:
    A, B -- Tuples representing the coordinates of points A and B (x, y).
    C -- Tuple representing the coordinates of the center point (x, y).
    length -- The length of the perpendicular line segment.

    Returns:
    a -- Slope of the perpendicular line (or `inf` for a vertical line).
    b -- Intercept of the perpendicular line (or x-coordinate if the line is vertical).
    start -- Starting coordinate of the perpendicular line segment.
    end -- Ending coordinate of the perpendicular line segment.
    """
    # Convert points to numpy arrays
    A, B, C = np.array(A), np.array(B), np.array(C)

    # Vector AB
    AB = B - A

    # Slope of AB
    slope_AB = (B[1] - A[1]) / (B[0] - A[0]) if (B[0] - A[0]) != 0 else np.inf

    # Slope of perpendicular line
    if slope_AB == 0:  # If AB is horizontal, perpendicular is vertical
        a = np.inf
    elif slope_AB == np.inf:  # If AB is vertical, perpendicular is horizontal
        a = 0
    else:
        a = -1 / slope_AB

    # Intercept of the perpendicular line
    if a != np.inf:
        b = C[1] - a * C[0]
    else:
        b = C[0]  # Vertical line, no y-intercept

    # Half-length of the line segment
    half_length = length / 2

    # Compute perpendicular direction vector
    if a == 0:  # Horizontal line
        start = C[0] - half_length
        end = C[0] + half_length
    elif a == np.inf:  # Vertical line
        start = C[1] - half_length
        end = C[1] + half_length
    else:
        perp_vec = np.array([1, a])  # Direction vector
        perp_vec /= np.linalg.norm(perp_vec)  # Normalize
        start = C[0] - half_length * perp_vec[0]
        end = C[0] + half_length * perp_vec[0]

    return a, b, start, end


def calculate_angle(A, B):
    """
    Calculate the angle (in degrees) between line segment AB and the horizontal axis.

    Parameters:
    A, B -- Tuples representing the coordinates of points A and B (x, y).

    Returns:
    angle -- Angle in degrees.
    """
    # Extract coordinates of points
    x1, y1 = A
    x2, y2 = B

    # Calculate the x and y components of vector AB
    dx = x2 - x1
    dy = y2 - y1

    # Compute slope
    if dx == 0:  # Prevent division by zero
        if dy > 0:
            angle = 90.0
        elif dy < 0:
            angle = -90.0
        else:
            angle = 0.0  # A and B are the same point
    else:
        angle = math.degrees(math.atan2(dy, dx))

    # Ensure the angle is in the range [0, 360)
    if angle < 0:
        angle += 360.0

    return angle



def create_sample_space(a_star_path, lane_width=2.75, step=25):
    """
    Create a sample space of perpendicular line segments along the given A* path.

    Parameters:
    a_star_path -- List of tuples representing the A* path coordinates (x, y).
    lane_width -- Width of the lane for defining the perpendicular segment (default: 2.75 meters).
    step -- Step size to determine the spacing of line segments along the path (default: 25).

    Returns:
    sample_space -- A list of dictionaries, each representing a perpendicular line segment with the following keys:
        'a' -- Slope of the perpendicular line.
        'b' -- Intercept of the perpendicular line.
        'st' -- Start coordinate of the perpendicular line segment.
        'ed' -- End coordinate of the perpendicular line segment.
        'th' -- Angle of the path segment (in radians).
        'id' -- Index of the line segment in the sample space.
    """
    sample_space = []
    for i in range(1, int(len(a_star_path) / step)):
        idx = step * i
        # Calculate the parameters of the perpendicular line segment
        a, b, start, end = perpendicular_line_segment_params(
            a_star_path[idx - 1], a_star_path[idx + 1], a_star_path[idx],
            lane_width * 2
        )
        # Calculate the angle of the path segment
        th = math.radians(calculate_angle(a_star_path[idx - 1], a_star_path[idx + 1]))
        # Create a dictionary with the segment parameters
        dic = {'a': a, 'b': b, 'st': start, 'ed': end, 'th': th, 'id': i}
        sample_space.append(dic)
    return sample_space


# def sample_waypoint(sample_space_dic):
#     """
#     Randomly sample a waypoint within the given sample space segment.

#     Parameters:
#     sample_space_dic -- Dictionary representing a segment in the sample space, 
#                         containing slope ('a'), intercept ('b'), start ('st'), end ('ed'),
#                         angle ('th'), and other attributes.

#     Returns:
#     x -- x-coordinate of the sampled waypoint.
#     y -- y-coordinate of the sampled waypoint.
#     th -- Perturbed angle of the waypoint (in radians).
#     k -- Randomly sampled curvature value.
#     """
#     # Handle vertical lines (infinite slope)
#     if sample_space_dic['a'] == np.inf:
#         x = sample_space_dic['b']
#         y = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
#     # Handle horizontal lines (zero slope)
#     elif sample_space_dic['a'] == 0:
#         x = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
#         y = sample_space_dic['b']
#     # Handle general case with finite slope
#     else:
#         x = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
#         y = sample_space_dic['a'] * x + sample_space_dic['b']

#     # Perturb the angle with a normal distribution
#     th = np.random.normal(sample_space_dic['th'], np.pi / 18)
#     # Generate a random curvature value
#     k = np.random.normal(0, 0.1)

#     return x, y, th, k
def sample_waypoint(sample_space_dic, num_samples=30):
    """
    Randomly sample a waypoint within the given sample space segment.

    Parameters:
    sample_space_dic -- Dictionary representing a segment in the sample space, 
                        containing slope ('a'), intercept ('b'), start ('st'), end ('ed'),
                        angle ('th'), and other attributes.

    Returns:
    x -- x-coordinate of the sampled waypoint.
    y -- y-coordinate of the sampled waypoint.
    th -- Perturbed angle of the waypoint (in radians).
    k -- Randomly sampled curvature value.
    waypoints -- List of all sampled waypoints along the line segment, where each waypoint is a tuple (x, y).
    """
    # Original logic for randomly sampling a waypoint
    if sample_space_dic['a'] == np.inf:  # Vertical line
        x = sample_space_dic['b']
        y = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
    elif sample_space_dic['a'] == 0:  # Horizontal line
        x = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
        y = sample_space_dic['b']
    else:  # General case
        x = np.random.uniform(sample_space_dic['st'], sample_space_dic['ed'])
        y = sample_space_dic['a'] * x + sample_space_dic['b']

    # Perturb the angle with a normal distribution
    th = np.random.normal(sample_space_dic['th'], np.pi / 18)
    # Generate a random curvature value
    k = np.random.normal(0, 0.1)

    # Additional logic for generating all waypoints along the line segment
    waypoints = []
    if sample_space_dic['a'] == np.inf:  # Vertical line
        waypoint_x = sample_space_dic['b']
        waypoint_ys = np.linspace(sample_space_dic['st'], sample_space_dic['ed'], num_samples)
        for waypoint_y in waypoint_ys:
            waypoints.append((waypoint_x, waypoint_y))  # Only x, y
    elif sample_space_dic['a'] == 0:  # Horizontal line
        waypoint_y = sample_space_dic['b']
        waypoint_xs = np.linspace(sample_space_dic['st'], sample_space_dic['ed'], num_samples)
        for waypoint_x in waypoint_xs:
            waypoints.append((waypoint_x, waypoint_y))  # Only x, y
    else:  # General case
        waypoint_xs = np.linspace(sample_space_dic['st'], sample_space_dic['ed'], num_samples)
        for waypoint_x in waypoint_xs:
            waypoint_y = sample_space_dic['a'] * waypoint_x + sample_space_dic['b']
            waypoints.append((waypoint_x, waypoint_y))  # Only x, y

    return x, y, th, k, waypoints



class Node:
    def __init__(self, position, parent=None, idx=None):
        self.position = position  # Node coordinates (x, y, theta)
        self.parent = parent  # Parent node
        self.cost = 0.0  # Cost to reach this node from the start node
        self.idx = idx


def sample_random_point(sample_space, itr, goal):
    """
    Sample a random point from the sample space or return the goal point.

    Parameters:
    sample_space -- List of dictionaries representing the sample space segments.
    itr -- Current iteration index.
    goal -- Tuple representing the goal point (x, y, th, k).

    Returns:
    x, y -- Coordinates of the sampled point.
    th -- Heading angle of the sampled point.
    k -- Curvature of the sampled point.
    idx -- Index of the sampled segment.
    """
    d = len(sample_space)
    idx = itr % d + 1  # Wrap-around index
    if idx == d:
        return goal[0], goal[1], goal[2], goal[3], idx, None
    else:
        x, y, th, k, waypoints = sample_waypoint(sample_space[idx - 1])
        return x, y, th, k, idx, waypoints


def find_nearest(tree, point):
    """
    Find the nearest node in the tree to a given point.

    Parameters:
    tree -- List of nodes in the tree.
    point -- Tuple representing the point (x, y, th, k, idx).

    Returns:
    nearest_node -- Node in the tree closest to the given point based on Euclidean distance.
    """
    min_distance = float('inf')
    nearest_node = None
    for node in tree:
        if node.idx == point[-1] - 1:  # Ensure nodes are from the same segment
            current_distance = np.linalg.norm(np.array(node.position[0:2]) - np.array(point[0:2]))
            if current_distance < min_distance:
                min_distance = current_distance
                nearest_node = node
    return nearest_node


def find_nearby_nodes(tree, new_node):
    """
    Find nodes in the tree that are adjacent to the given node by segment index.

    Parameters:
    tree -- List of nodes in the tree.
    new_node -- Node for which to find nearby nodes.

    Returns:
    nearby_nodes -- List of nodes that are adjacent to the new_node in terms of segment index.
    """
    nearby_nodes = []
    for node in tree:
        if node.idx == new_node.idx - 1 or node.idx == new_node.idx + 1:
            nearby_nodes.append(node)
    return nearby_nodes


def steer(q_nearest, q_rand, glb_map):
    """
    Generate a path from the nearest node to the random point, checking for boundary and collision constraints.

    Parameters:
    q_nearest -- Node representing the nearest point in the tree.
    q_rand -- Tuple representing the randomly sampled point (x, y, th, k).
    glb_map -- Dictionary representing the global map with boundaries and obstacles.

    Returns:
    new_node -- Node representing the new position, or None if the path is invalid.
    path_x, path_y -- Coordinates of the generated path, or None if invalid.
    """
    # Generate a G2 clothoid curve path
    path_x, path_y, _, path_length, _ = clothoid_path(
        q_nearest.position[0], q_nearest.position[1], q_nearest.position[2],
        q_nearest.position[3], q_rand[0], q_rand[1], q_rand[2], q_rand[3]
    )

    # Check if the path stays within map boundaries
    for x, y in zip(path_x, path_y):
        if (x <= -glb_map['l'] / 2 + glb_map['m'] or x >= glb_map['l'] / 2 - glb_map['m']
                or y <= -glb_map['w'] / 2 + glb_map['m'] or y >= glb_map['w'] / 2 - glb_map['m']):
            return None, None, None

    # Check for collisions with obstacles
    for ob in glb_map['o']:
        if ob['type'] == 'box':
            low_bound = (ob['geometry']['position'][0] - ob['geometry']['length'] / 2 - glb_map['m'],
                         ob['geometry']['position'][1] - ob['geometry']['width'] / 2 - glb_map['m'])
            up_bound = (ob['geometry']['position'][0] + ob['geometry']['length'] / 2 + glb_map['m'],
                        ob['geometry']['position'][1] + ob['geometry']['width'] / 2 + glb_map['m'])
            for x, y in zip(path_x, path_y):
                if low_bound[0] <= x <= up_bound[0] and low_bound[1] <= y <= up_bound[1]:
                    return None, None, None
        elif ob['type'] == 'sphere':
            center = (ob['geometry']['position'][0], ob['geometry']['position'][1])
            for px, py in zip(path_x, path_y):
                if np.linalg.norm(np.array(center) - np.array([px, py])) <= ob['geometry']['radius'] + glb_map['m']:
                    return None, None, None

    # Create a new node and update its cost
    new_node = Node(position=np.array(q_rand[0:4]), parent=q_nearest, idx=q_rand[-1])
    new_node.cost = q_nearest.cost + path_length
    return new_node, path_x, path_y



def rewire(new_node, nearby_nodes, glb_map):
    """
    Rewire nearby nodes to the new node if a shorter path is found.

    Parameters:
    new_node -- The newly added node in the tree.
    nearby_nodes -- List of nodes near the new node.
    glb_map -- Dictionary representing the global map, including boundaries and obstacles.

    Returns:
    None. Updates the parent and cost of nearby nodes if rewiring occurs.
    """
    for node in nearby_nodes:
        # Generate a G2 clothoid path from the new node to the nearby node
        path_x, path_y, _, path_length, _ = clothoid_path(
            new_node.position[0], new_node.position[1], new_node.position[2], new_node.position[3],
            node.position[0], node.position[1], node.position[2], node.position[3]
        )

        # Check for collisions and update parent if a shorter path is found
        if is_collision_free(path_x, path_y, glb_map):
            new_cost = new_node.cost + path_length
            if new_cost < node.cost:
                node.parent = new_node
                node.cost = new_cost


def is_collision_free(path_x, path_y, glb_map):
    """
    Check if a given path is collision-free within the global map.

    Parameters:
    path_x, path_y -- Lists of x and y coordinates of the path.
    glb_map -- Dictionary representing the global map, including boundaries and obstacles.

    Returns:
    True if the path is collision-free, False otherwise.
    """
    for x, y in zip(path_x, path_y):
        # Check boundaries
        if (x <= -glb_map['l'] / 2 + glb_map['m'] or x >= glb_map['l'] / 2 - glb_map['m']
                or y <= -glb_map['w'] / 2 + glb_map['m'] or y >= glb_map['w'] / 2 - glb_map['m']):
            return False

        # Check collisions with obstacles
        for ob in glb_map['o']:
            if ob['type'] == 'box':
                low_bound = (ob['geometry']['position'][0] - ob['geometry']['length'] / 2 - glb_map['m'],
                             ob['geometry']['position'][1] - ob['geometry']['width'] / 2 - glb_map['m'])
                up_bound = (ob['geometry']['position'][0] + ob['geometry']['length'] / 2 + glb_map['m'],
                            ob['geometry']['position'][1] + ob['geometry']['width'] / 2 + glb_map['m'])
                if low_bound[0] <= x <= up_bound[0] and low_bound[1] <= y <= up_bound[1]:
                    return False
            elif ob['type'] == 'sphere':
                center = (ob['geometry']['position'][0], ob['geometry']['position'][1])
                if np.linalg.norm(np.array(center) - np.array([x, y])) <= ob['geometry']['radius'] + glb_map['m']:
                    return False
    return True


def reachable_point_to_line(point, wayline_dic, glb_map):
    """
    Check if a point can reach a line segment within the global map.

    Parameters:
    point -- Starting point (x, y, th, k).
    wayline_dic -- Dictionary representing a line segment in the sample space.
    glb_map -- Dictionary representing the global map.

    Returns:
    True if reachable, False otherwise.
    """
    threshold = 0
    while threshold < 3:
        waypoint = sample_random_point([wayline_dic, wayline_dic], 0, [0, 0, 0, 0])
        new_node, _, _ = steer(Node(position=point), waypoint, glb_map)
        if new_node is None:
            threshold += 1
        else:
            return True
    return False


def reachable_line_to_point(wayline_dic, point, glb_map):
    """
    Check if a line segment can reach a point within the global map.

    Parameters:
    wayline_dic -- Dictionary representing a line segment in the sample space.
    point -- Target point (x, y, th, k).
    glb_map -- Dictionary representing the global map.

    Returns:
    True if reachable, False otherwise.
    """
    threshold = 0
    while threshold < 3:
        waypoint = sample_random_point([wayline_dic, wayline_dic], 0, [0, 0, 0, 0])
        new_node, _, _ = steer(Node(position=waypoint), point, glb_map)
        if new_node is None:
            threshold += 1
        else:
            return True
    return False


def reachable_line_to_line(wayline_dic1, wayline_dic2, glb_map):
    """
    Check if one line segment can reach another within the global map.

    Parameters:
    wayline_dic1, wayline_dic2 -- Dictionaries representing two line segments in the sample space.
    glb_map -- Dictionary representing the global map.

    Returns:
    True if reachable, False otherwise.
    """
    threshold = 0
    while threshold < 3:
        random_point = sample_waypoint(wayline_dic1)
        if reachable_point_to_line(random_point, wayline_dic2, glb_map):
            return True
        else:
            threshold += 1
    return False


def reachable(i, j, sample_space, glb_map, start, goal):
    """
    Determine reachability between nodes or points within the sample space.

    Parameters:
    i, j -- Indices or special strings ('start', 'goal') representing the source and target.
    sample_space -- List of dictionaries representing the sample space.
    glb_map -- Dictionary representing the global map.
    start, goal -- Points representing the start and goal positions.

    Returns:
    True if reachable, False otherwise.
    """
    if i == 'start' and j == 'goal':
        var = False
    elif i == 'start' and j != 'goal':
        var = reachable_point_to_line(start, sample_space[j - 1], glb_map)
    elif i != 'start' and j != 'goal':
        var = reachable_line_to_line(sample_space[i], sample_space[j], glb_map)
    else:
        var = reachable_line_to_point(sample_space[i], goal, glb_map)

    return var


def refinement(sample_space, glb_map, start, goal):
    """
    Refine a path by removing unnecessary intermediate nodes.

    Parameters:
    sample_space -- List of dictionaries representing the sample space.
    glb_map -- Dictionary representing the global map.
    start, goal -- Points representing the start and goal positions.

    Returns:
    optimized_path -- List of indices or special strings ('start', 'goal') representing the optimized path.
    """
    path = ['start']
    for i in range(1, len(sample_space)):
        path.append(i)
    path.append('goal')

    optimized_path = []
    ii = 0
    while ii < len(path) - 1:
        for jj in range(len(path) - 1, ii, -1):
            if reachable(path[ii], path[jj], sample_space, glb_map, start, goal):
                optimized_path.append(path[ii])
                ii = jj
                break
            if jj == ii + 1:
                optimized_path.append(path[ii + 1])
                ii += 1
    return optimized_path



def rrt_star(glb_map, a_star_path, start, goal, iterations=200, max_iterations=300, save_pkl = True, print_log = True):
    """
    Perform the RRT* (Rapidly-exploring Random Tree Star) algorithm to find a path from start to goal.

    Parameters:
    glb_map -- Dictionary representing the global map, including boundaries and obstacles.
    a_star_path -- List of coordinates representing the initial A* path.
    start -- Starting position (x, y, th, k).
    goal -- Goal position (x, y, th, k).
    iterations -- Number of iterations to attempt to find a solution (default: 100).
    max_iterations -- Maximum number of iterations allowed to prevent infinite loops (default: 200).
    save_pkl -- Save path to a pkl file
    print_log -- Print debug information

    Returns:
    path -- List of refined nodes representing the final path.
    path_x, path_y -- Lists of x and y coordinates of the explored points.
    found_path_x, found_path_y -- Lists of x and y coordinates of the final found path.
    found_path_yaw -- List of yaw angles along the final found path.
    lgt -- Total length of the final found path.
    found_path_curvature -- List of curvatures along the final found path.
    waypoints_list -- Dictionary of waypoints for each idx sampled during the RRT process.
    """
    start_node = Node(position=start, idx=0)
    goal_node = Node(position=goal)
    goal_node.cost = float('inf')
    tree = [start_node]
    path_x, path_y = [], []
    waypoints_list = {}  # Initialize as a dictionary
    sample_space = create_sample_space(a_star_path)

    iteration = 0
    max_iteration = 0
    backward = 0

    while iteration < iterations and max_iteration < max_iterations:
        max_iteration += 1
        if print_log:
            print(f"---------iteration={iteration}---------")

        # Sample a random point or goal
        result = sample_random_point(sample_space, iteration, goal)  # 调用一次
        random_point = result[:5]  # 前 5 个值
        waypoints = result[5]      # 第 6 个值 
        idx = random_point[-1]  # Use idx from random_point

        if waypoints:
            waypoints_list[idx] = waypoints  # Only keep the latest waypoints for each idx
        if print_log:
            print(f"Sampled point: {random_point}")

        # Find the nearest node in the tree and attempt to steer toward the random point
        nearest_node = find_nearest(tree, random_point)
        new_node, x, y = steer(nearest_node, random_point, glb_map)

        # If the path to the new node is invalid, handle backward movement
        if new_node is None:
            if print_log:
                print("Path invalid, retrying...")
            backward += 1
            if backward >= 5 and iteration >= 1:
                iteration -= 1
                backward = 0
            continue

        # Add the new node and its path to the tree
        tree.append(new_node)
        path_x.extend(x)
        path_y.extend(y)
        iteration += 1

        # Rewire nearby nodes to potentially optimize the tree
        nearby_nodes = find_nearby_nodes(tree, new_node)
        rewire(new_node, nearby_nodes, glb_map)

        # Check if the goal is reached and update goal node cost if necessary
        if (new_node.position[0], new_node.position[1], new_node.position[2]) == (goal[0], goal[1], goal[2]):
            if print_log:
                print(f"Goal reached at position: {new_node.position}")
            goal_node_cost = new_node.cost
            if goal_node_cost < goal_node.cost:
                goal_node.parent = new_node.parent
                goal_node.cost = goal_node_cost
                goal_node.idx = new_node.idx

    # Construct the path from the goal to the start
    current_node = goal_node
    node_list = []
    node_position_list = []
    while current_node:
        node_list.append(current_node)
        node_position_list.append(current_node.position)
        current_node = current_node.parent
    node_list.reverse()
    node_position_list.reverse()

    # Refine the path to reduce unnecessary nodes
    path = []
    q = 0
    while q < len(node_list) - 1:
        for p in range(len(node_list) - 1, q, -1):
            new, _, _ = steer(node_list[q], node_list[p].position, glb_map)
            if new:
                path.append(node_list[q].position)
                q = p
                break
    if path:
        path.append(goal)
    if not path:
        # Fallback to direct node connections if refinement fails
        current_node = goal_node
        while current_node:
            path.append(current_node.position)
            current_node = current_node.parent
        path.reverse()

    # Generate the detailed path with clothoid curves
    found_path_x, found_path_y, found_path_yaw, found_path_curvature = [], [], [], []
    lgt = 0
    for i in range(len(path) - 1):
        found_x, found_y, found_yaw, lg, found_curvature = clothoid_path(
            path[i][0], path[i][1], path[i][2], path[i][3],
            path[i + 1][0], path[i + 1][1], path[i + 1][2], path[i + 1][3]
        )
        found_path_x.extend(found_x)
        found_path_y.extend(found_y)
        found_path_yaw.extend(found_yaw)
        found_path_curvature.extend(found_curvature)
        lgt += lg
    
    import pandas as pd
    import time
    if save_pkl:
        data = {
            'origin_nodes_list': pd.Series(node_position_list),
            'refined_nodes_list': pd.Series(path),
            'explored_x': pd.Series(path_x),
            'explored_y': pd.Series(path_y),
            'found_x': pd.Series(found_path_x),
            'found_y': pd.Series(found_path_y),
        }
        df = pd.DataFrame(data)
        df.to_pickle(f"{time.strftime('%Y-%m-%d_%H_%M_%S', time.localtime())}.pkl")     

    return path, path_x, path_y, found_path_x, found_path_y, found_path_yaw, lgt, found_path_curvature, waypoints_list



def draw_path_rrt(ax, path, path_x, path_y, found_path_x, found_path_y):
    """
    Visualize the RRT* algorithm results.

    Parameters:
    ax -- Matplotlib axis object for plotting.
    path -- List of nodes representing the refined path.
    path_x, path_y -- Lists of x and y coordinates of explored points.
    found_path_x, found_path_y -- Lists of x and y coordinates of the final found path.

    Returns:
    None. The plot is updated.
    """
    if path is None:
        print("No RRT Star path found!")
        return
    ax.scatter(path_x, path_y, color="#48597e", marker="o", s=0.5, label="Explored")
    ax.scatter(found_path_x, found_path_y, color="red", marker="o", s=1, label="Found Path")
    node_x = [p[0] for p in path]
    node_y = [p[1] for p in path]
    ax.scatter(node_x, node_y, color="blue", marker="x", s=100, label="Nodes")
    ax.legend()
    plt.draw()

def compare_paths(ax, original_path_x, original_path_y, refined_path_x, refined_path_y, save_folder="compare_path"):
    """
    Compare and visualize the original and refined paths on an existing plot,
    and save the plot to a specified folder.

    Parameters:
    ax -- Matplotlib axis object for plotting.
    original_path_x, original_path_y -- Coordinates of the original path (lists of x and y points).
    refined_path_x, refined_path_y -- Coordinates of the refined path (lists of x and y points).
    save_folder -- Folder to save the plot (default: "compare_path").

    Returns:
    None. The plot is updated on the given axis and saved as a PNG file.
    """
    # Plot the original path with dotted line and 'o' markers
    if original_path_x and original_path_y:
        ax.scatter(original_path_x, original_path_y, color='blue', s=1, marker='o', label="Original Path Points")
    else:
        print("No Original Path data provided!")

    # Plot the refined path with solid line and 'x' markers
    if refined_path_x and refined_path_y:
        ax.scatter(refined_path_x, refined_path_y, color='red', s=0.5, marker='x', label="Refined Path Points")
    else:
        print("No Refined Path data provided!")

    # Add a legend
    ax.legend(fontsize=12)

    # Ensure the save folder exists
    os.makedirs(save_folder, exist_ok=True)

    # Generate a timestamped file name
    timestamp = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
    save_path = os.path.join(save_folder, f"compare_paths_{timestamp}.png")

    # Save the plot
    plt.savefig(save_path, format='png', dpi=300)
    print(f"Plot saved as: {os.path.abspath(save_path)}")

    # Show the plot
    plt.show()

def plot_waypoints(waypoints_dict, ax, title="Waypoints Scatter Plot", color='green', size=5):
    """
    Plot all waypoints from a dictionary on a 2D scatter plot using a given Axes object,
    and save the plot to the 'compare_LTR' folder in the current directory.

    Parameters:
    waypoints_dict -- Dictionary of waypoints where each key maps to a list of tuples (x, y).
    ax -- Matplotlib Axes object where the scatter plot will be drawn.
    title -- Title of the plot (default: "Waypoints Scatter Plot").
    color -- Color of the scatter points (default: 'blue').
    size -- Size of the scatter points (default: 20).

    Returns:
    None. Modifies the provided Axes object and saves the plot as a PNG file.
    """
    # Validate input data
    if not isinstance(waypoints_dict, dict):
        raise ValueError("waypoints_dict must be a dictionary")
    for key, value in waypoints_dict.items():
        if not isinstance(value, list) or not all(isinstance(coord, tuple) and len(coord) == 2 for coord in value):
            raise ValueError("Each value in waypoints_dict must be a list of (x, y) tuples")

    # Collect all x and y coordinates from the dictionary
    waypoint_x = []
    waypoint_y = []
    for waypoints in waypoints_dict.values():
        for wp in waypoints:
            waypoint_x.append(wp[0])
            waypoint_y.append(wp[1])

    # Create the scatter plot
    ax.scatter(waypoint_x, waypoint_y, c=color, s=size, label='Waypoints')

    # Add title and axis labels
    ax.set_title(title, fontsize=16)
    ax.set_xlabel("X Coordinate", fontsize=12)
    ax.set_ylabel("Y Coordinate", fontsize=12)

    # Add grid and legend
    ax.grid(True, linestyle='--', alpha=0.7)
    ax.legend()

    # Create the save folder if it does not exist
    save_folder = "waypoints_img"
    os.makedirs(save_folder, exist_ok=True)

    # Generate a timestamped file name
    timestamp = time.strftime("%Y-%m-%d_%H-%M-%S", time.localtime())
    save_path = os.path.join(save_folder, f"waypoints_plot_{timestamp}.png")

    # Save the plot
    plt.savefig(save_path, format='png', dpi=300)
    print(f"Plot saved as: {os.path.abspath(save_path)}")

    # Show the plot
    plt.show()